#include "G4SystemOfUnits.hh"
#include "Randomize.hh"
#include "G4ParticleTable.hh"
#include "G4ParticleDefinition.hh"
#include <memory>

#include "MyGenCreator.hh"
#include "MyGenContainer.hh"

#include "MarvelPhoton.hh"

MarvelPhoton::MarvelPhoton(): MyBaseGen()
{
    m_PGType = "Simple Particle Gun";
    m_ClassName = "MarvelPhoton";
    m_TypeFlag = 1001;
    m_EvtFlag = 0;

    m_Particle = "gamma";
    m_ParticleEnergy = 1. * MeV;
    m_ParticlePosition = G4ThreeVector(0., 0., 0.);
    m_ParticlePolarization = G4ThreeVector(0., 0., 0.);
    m_ParticleMomentumDirection = G4ThreeVector(0., 0., 0.);

    m_Gun = nullptr;
    m_PosGen = nullptr;
}

MarvelPhoton::~MarvelPhoton()
{
    if(!m_Gun)
        delete m_Gun;
    if(!m_FileGammaE)
        delete m_FileGammaE;
}

void MarvelPhoton::Initialize(YAML::Node yamlNode)
{
    m_Gun = new G4ParticleGun(1);
    m_PosGen = MyGenContainer::GetInstance().GetMyPos();

    m_FileGammaE = new TFile(yamlNode["HistFile"].as<std::string>("GammaEnergy.root").c_str());
    m_HistGammaE = dynamic_cast<TH1F*>(m_FileGammaE->Get(yamlNode["HistName"].as<std::string>("GammaEnergy").c_str()));

    G4double loE = yamlNode["LowEnergy"].as<double>(0.) * MeV;
    G4double hiE = yamlNode["HighEnergy"].as<double>(100.) * MeV;

    G4int loBin = m_HistGammaE->FindBin(loE);
    G4int hiBin = m_HistGammaE->FindBin(hiE);
    G4int totalBin = m_HistGammaE->GetNbinsX();

    if(loBin == 0) loBin = 1;
    if(hiBin >= totalBin) hiBin = totalBin;
    if(loBin >= hiBin)
    {
        std::cerr << "Error! LowEnergy larger than HighEnergy" << std::endl;
        exit(-1);
    }
    if(loBin != 1 || hiBin != totalBin)
    {
        TH1F* newHist = new TH1F("newGammaEnergy", "", hiBin - loBin + 1, m_HistGammaE->GetBinLowEdge(loBin), m_HistGammaE->GetBinLowEdge(hiBin+1));
        for(int ib = 0; ib < hiBin - loBin + 1; ++ib)
        {
            newHist->SetBinContent(ib+1, m_HistGammaE->GetBinContent(ib+loBin));
        }
        m_HistGammaE = newHist;
    }

    G4ParticleTable *particleTable = G4ParticleTable::GetParticleTable();
    G4ParticleDefinition *particle = particleTable->FindParticle(m_Particle);
    m_Gun->SetParticleDefinition(particle);
}

void MarvelPhoton::GeneratePrimaryVertex(G4Event* anEvent)
{
    m_ParticleEnergy = m_HistGammaE->GetRandom() * MeV;
    m_Gun->SetParticleEnergy(m_ParticleEnergy);

    PosAndMom PosAndMomDir = m_PosGen->GenPosAndMomDir();
    m_ParticlePosition = PosAndMomDir.first;
    m_ParticleMomentumDirection = PosAndMomDir.second;
    m_Gun->SetParticlePosition(m_ParticlePosition);
    m_Gun->SetParticleMomentumDirection(m_ParticleMomentumDirection);

    m_Gun->GeneratePrimaryVertex(anEvent);
    ++m_EvtFlag;
}

REGISTGPSGEN(MarvelPhoton)